#ifndef _TLG_H_
#define _TLG_H_

#include <stdint.h>

uint32_t *tlg_load(uint8_t *buf, uint64_t len, int *w, int *h);

#ifdef TLG_IMPLEMENTATION

typedef struct {
    uint8_t channels;
    uint32_t width;
    uint32_t height;
    uint32_t blockh;
} __attribute__((packed)) tlg5_t;

typedef struct {
    uint8_t channels;
    uint8_t flags;
    uint8_t type;
    uint8_t external_table;
    uint32_t width;
    uint32_t height;
    uint32_t maxbitlen;
} __attribute__((packed)) tlg6_t;

typedef struct {
    uint8_t text[4096];
    uint16_t offset;
} lzss_t;
short golomb_compression_table[4][9] = { {3, 7, 15, 27, 63, 108, 223, 448, 130 }, {3, 5, 13, 24, 51, 95, 192, 384, 257 },
    {2, 5, 12, 21, 39, 86, 155, 320, 384 }, {2, 3, 9, 18, 33, 61, 129, 258, 511 } };
uint8_t leading_zero_table[4096];
uint8_t golomb_bit_length_table[4 * 2 * 128][4];

void lzss_decompress(lzss_t *ctx, uint8_t *input, uint64_t input_size, uint8_t *output)
{
    uint8_t *end = input + input_size, x0, x1, c;
    int j, flags = 0, position, length;

    memset(output, 0, input_size);

    while(input < end) {
        flags >>= 1;
        if((flags & 0x100) != 0x100)
            flags = *input++ | 0xff00;

        if((flags & 1) == 1) {
            x0 = *input++;
            x1 = *input++;
            position = x0 | ((x1 & 0xf) << 8);
            length = 3 + ((x1 & 0xf0) >> 4);
            if(length == 18)
                length += *input++;
            for(j = 0; j < length; j ++) {
                c = ctx->text[position];
                *output++ = c;
                ctx->text[ctx->offset] = c;
                ctx->offset++;
                ctx->offset &= 0xfff;
                position++;
                position &= 0xfff;
            }
        } else {
            c = *input++;
            *output++ = c;
            ctx->text[ctx->offset] = c;
            ctx->offset++;
            ctx->offset &= 0xfff;
        }
    }
}

__inline__ uint32_t tlg6_make_gt_mask(uint32_t a, uint32_t b)
{
    uint32_t tmp2 = ~b;
    uint32_t tmp = ((a & tmp2) + (((a ^ tmp2) >> 1) & 0x7f7f7f7f)) & 0x80808080;
    return ((tmp >> 7) + 0x7f7f7f7f) ^ 0x7f7f7f7f;
}

__inline__ uint32_t tlg6_packed_bytes_add(uint32_t a, uint32_t b)
{
    return a + b - ((((a & b) << 1) + ((a ^ b) & 0xfefefefe)) & 0x01010100);
}

__inline__ uint32_t tlg6_med(uint32_t a, uint32_t b, uint32_t c, uint32_t v)
{
    uint32_t aa_gt_bb = tlg6_make_gt_mask(a, b);
    uint32_t a_xor_b_and_aa_gt_bb = ((a ^ b) & aa_gt_bb);
    uint32_t aa = a_xor_b_and_aa_gt_bb ^ a;
    uint32_t bb = a_xor_b_and_aa_gt_bb ^ b;
    uint32_t n = tlg6_make_gt_mask(c, bb);
    uint32_t nn = tlg6_make_gt_mask(aa, c);
    uint32_t m = ~(n | nn);
    return tlg6_packed_bytes_add((n & aa) | (nn & bb) | ((bb & m) - (c & m) + (aa & m)), v);
}

__inline__ uint32_t tlg6_avg(uint32_t a, uint32_t b, uint32_t c, uint32_t v)
{
    (void)c;
    return tlg6_packed_bytes_add((a & b) + (((a ^ b) & 0xfefefefe) >> 1) + ((a ^ b) & 0x01010101), v);
}

void tlg6_decode_line(uint32_t *prev_line, uint32_t *current_line, int width, int start_block, int block_limit,
    uint8_t *filter_types, int skip_block_bytes, uint32_t *in, uint32_t initialp, int odd_skip, int dir)
{
    uint8_t a, r, g, b;
    uint32_t u, p, up;
    int step, i, w, ww;

    if(start_block) {
        prev_line += start_block << 3;
        current_line += start_block << 3;
        p = current_line[-1];
        up = prev_line[-1];
    } else
        p = up = initialp;

    in += skip_block_bytes * start_block;
    step = (dir & 1) ? 1 : -1;
    for(i = start_block; i < block_limit; i++) {
        w = width - (i << 3);
        if(w > 8) w = 8;
        ww = w; if(step == -1) in += ww - 1;
        if(i & 1) in += odd_skip * ww;
        do {
            a = (*in >> 24) & 0xff;
            r = (*in >> 16) & 0xff;
            g = (*in >> 8) & 0xff;
            b = (*in) & 0xff;
            switch(filter_types[i] >> 1) {
                case 1: r += g; b += g; break;
                case 2: g += b; r += g; break;
                case 3: g += r; b += g; break;
                case 4: b += r; g += b; r += g; break;
                case 5: b += r; g += b; break;
                case 6: b += g; break;
                case 7: g += b; break;
                case 8: r += g; break;
                case 9: r += b; g += r; b += g; break;
                case 10:b += r; g += r; break;
                case 11: r += b; g += b; break;
                case 12: r += b; g += r; break;
                case 13: b += g; r += b; g += r; break;
                case 14: g += r; b += g; r += b; break;
                case 15: g += (b << 1); r += (b << 1); break;
                default: break;
            }
            u = *prev_line;
            p = filter_types[i] & 1 ?
                tlg6_avg(p, u, up, (0xff0000 & (r << 16)) + (0xff00 & (g << 8)) + (0xff & b) + (a << 24)) :
                tlg6_med(p, u, up, (0xff0000 & (r << 16)) + (0xff00 & (g << 8)) + (0xff & b) + (a << 24));
            p |= initialp;
            up = u;
            *current_line = p;
            current_line++;
            prev_line++;
            in += step;
        } while(-- w);
        in += skip_block_bytes + (step == 1 ? -ww : 1);
        if(i & 1) in -= odd_skip * ww;
    }
}

void tlg6_decode_golomb_values(uint8_t *pixel_buf, int pixel_count, uint8_t *bit_pool)
{
    int n = 3, a = 0, bit_pos = 1, count, b, bit_count, k, v, sign;
    uint8_t zero = (*bit_pool & 1) ? 0 : 1, *limit = pixel_buf + pixel_count * 4;
    uint32_t t;

    while(pixel_buf < limit) {
        t = *((uint32_t*)(bit_pool)) >> bit_pos;
        b = leading_zero_table[t & 4095];
        bit_count = b;
        while(!b) {
            bit_count += 12;
            bit_pos += 12;
            bit_pool += bit_pos >> 3;
            bit_pos &= 7;
            t = *(uint32_t*)(bit_pool) >> bit_pos;
            b = leading_zero_table[t & 4095];
            bit_count += b;
        }
        bit_pos += b;
        bit_pool += bit_pos >> 3;
        bit_pos &= 7;
        bit_count --;
        count = 1 << bit_count;
        t = *(uint32_t*)(bit_pool);
        count += ((t >> bit_pos) & (count - 1));
        bit_pos += bit_count;
        bit_pool += bit_pos >> 3;
        bit_pos &= 7;
        if(zero) {
            do { *pixel_buf = 0; pixel_buf += 4; } while(--count);
        } else {
            do {
                if((t = *(uint32_t*)(bit_pool) >> bit_pos)) {
                    b = leading_zero_table[t & 4095];
                    bit_count = b;
                    while(!b) {
                        bit_count += 12;
                        bit_pos += 12;
                        bit_pool += bit_pos >> 3;
                        bit_pos &= 7;
                        t = *(uint32_t*)(bit_pool) >> bit_pos;
                        b = leading_zero_table[t & 4095];
                        bit_count += b;
                    }
                    bit_count--;
                } else {
                    bit_pool += 5;
                    bit_count = bit_pool[-1];
                    bit_pos = 0;
                    t = *(uint32_t*)(bit_pool);
                    b = 0;
                }
                k = golomb_bit_length_table[a][n];
                v = (bit_count << k) + ((t >> b) & ((1 << k) - 1));
                sign = (v & 1) - 1;
                v >>= 1;
                a += v;
                *((uint8_t*)pixel_buf) = ((v ^ sign) + sign + 1);
                pixel_buf += 4;
                bit_pos += b;
                bit_pos += k;
                bit_pool += bit_pos >> 3;
                bit_pos &= 7;
                if(--n < 0) { a >>= 1; n = 3; }
            } while(--count);
        }
        zero ^= 1;
    }
}

uint32_t *tlg_load(uint8_t *buf, uint64_t len, int *w, int *h)
{
    lzss_t lzss;
    tlg5_t *hdr5;
    tlg6_t *hdr6;
    uint8_t version, *blkdata[4] = { 0 }, channel, mark, prevr, prevg, prevb, preva, outr, outg, outb, outa, r, g, b, a;
    uint8_t *ftypes = NULL, *ptr;
    uint32_t *ret = NULL, blksize, Y, y, ylim, x, above, x_block_count, y_block_count, *zero_line = NULL, *prev_line, *curr_line;
    uint32_t *pixbuf = NULL, main_count;
    int i, j, k, l, skip, fraction, pixel_count, block_y_shift, prev_y_shift, y_shift;

    if(!buf || len < 32 || !w || !h) return NULL;
    if(!memcmp(buf, "TLG0.0\0sds\x1a", 11)) buf += 15;
    if(memcmp(buf + 6, "\0raw\x1a", 5) || (buf[11] != 3 && buf[11] != 4)) return NULL;
    if(!memcmp(buf, "TLG5.0", 6)) { version = 5; } else
    if(!memcmp(buf, "TLG6.0", 6)) { version = 6; } else
    if(!memcmp(buf, "XXXYYY", 6)) { version = 5; buf[0x0C] ^= 0xAB; buf[0x10] ^= 0xAC; } else
    if(!memcmp(buf, "XXXZZZ", 6)) { version = 6; buf[0x0F] ^= 0xAB; buf[0x13] ^= 0xAC; } else
    if(!memcmp(buf, "JKMXE8", 6)) { version = 5; buf[0x0C] ^= 0x1A; buf[0x10] ^= 0x1C; } else return NULL;
    buf += 11;

    if(version == 5) {
        hdr5 = (tlg5_t*)buf; buf += sizeof(tlg5_t);
        if(!hdr5->blockh || hdr5->width > 4096 || hdr5->height > 4096) return NULL;
        *w = hdr5->width; *h = hdr5->height;
        if((ret = (uint32_t*)malloc(hdr5->width * hdr5->height * 4))) {
            for(channel = 0; channel < hdr5->channels; channel++)
                if(!(blkdata[channel] = (uint8_t*)malloc(hdr5->width * hdr5->height))) { free(ret); ret = NULL; goto cleanup; }
            buf += ((hdr5->height - 1) / hdr5->blockh + 1) * 4;
            memset(&lzss, 0, sizeof(lzss));
            for(Y = 0; Y < hdr5->height; Y += hdr5->blockh) {
                for(channel = 0; channel < hdr5->channels; channel++) {
                    mark = *buf++;
                    memcpy(&blksize, buf, 4); buf += 4;
                    if(!mark) lzss_decompress(&lzss, buf, blksize, blkdata[channel]);
                    else memcpy(blkdata[channel], buf, blksize);
                    buf += blksize;
                }
                ylim = Y + hdr5->blockh;
                if(ylim > hdr5->height) ylim = hdr5->height;
                for(y = Y; y < ylim; y++) {
                    prevr = prevg = prevb = preva = 0;
                    block_y_shift = (y - Y) * hdr5->width;
                    prev_y_shift = (y - 1) * hdr5->width;
                    y_shift = y * hdr5->width;
                    for(x = 0; x < hdr5->width; x++) {
                        a = !blkdata[3] ? 0 : blkdata[2][block_y_shift + x];
                        r = blkdata[2][block_y_shift + x];
                        g = blkdata[1][block_y_shift + x];
                        b = blkdata[0][block_y_shift + x];
                        r += g;
                        b += g;
                        prevr += r; outr = prevr;
                        prevg += g; outg = prevg;
                        prevb += b; outb = prevb;
                        preva += a; outa = preva;
                        if(y > 0) {
                            above = ret[prev_y_shift + x];
                            outr += (above >> 16) & 0xff;
                            outg += (above >> 8) & 0xff;
                            outb += (above) & 0xff;
                            outa += (above >> 24) & 0xff;
                        }
                        if(!blkdata[3]) outa = 0xff;
                        ret[y_shift + x] = (outb) | (outg << 8) | (outr << 16) | (outa << 24);
                    }
                }
            }
        }
    } else {
        hdr6 = (tlg6_t*)buf; buf += sizeof(tlg6_t);
        if(hdr6->width > 4096 || hdr6->height > 4096) return NULL;
        *w = hdr6->width; *h = hdr6->height;
        if((ret = (uint32_t*)malloc(hdr6->width * hdr6->height * 4))) {
            memset(ret, 0, hdr6->width * hdr6->height * 4);
            if(!(pixbuf = (uint32_t*)malloc(hdr6->width * hdr6->height * 16)) ||
                !(zero_line = (uint32_t*)malloc(hdr6->width * 4))) { free(ret); ret = NULL; goto cleanup; }
            memset(pixbuf, 0, hdr6->width * hdr6->height * 16);
            memset(zero_line, 0, hdr6->width * 4);
            memcpy(&blksize, buf, 4); buf += 4;
            x_block_count = ((hdr6->width - 1) >> 3) + 1;
            y_block_count = ((hdr6->height - 1) >> 3) + 1;
            x = x_block_count * y_block_count;
            if(!(ftypes = (uint8_t*)malloc(x))) { free(ret); ret = NULL; goto cleanup; }
            memset(ftypes, 0, x);
            memset(&lzss, 0, sizeof(lzss));
            ptr = lzss.text;
            for(i = 0; i < 32; i++) {
                for(j = 0; j < 16; j++) {
                    for(k = 0; k < 4; k++) *ptr++ = i;
                    for(k = 0; k < 4; k++) *ptr++ = j;
                }
            }
            lzss_decompress(&lzss, buf, blksize, ftypes); buf += blksize;
            prev_line = zero_line;
            for(i = 0; i < 4096; i++) {
                for(k = 0, j = 1; j != 4096 && !(i & j); j <<= 1, k++);
                k++; if(j == 4096) k = 0;
                leading_zero_table[i] = k;
            }
            for(k = 0; k < 4; k++) {
                for(l = i = 0; i < 9; i++) {
                    for(j = 0; j < golomb_compression_table[k][i]; j ++)
                        golomb_bit_length_table[l++][k] = i;
                }
            }

            main_count = hdr6->width >> 3;
            fraction = hdr6->width - (main_count << 3);
            for(y = 0; y < hdr6->height; y += 8) {
                ylim = y + 8; if(ylim >= hdr6->height) ylim = hdr6->height;
                pixel_count = (ylim - y) * hdr6->width;
                for(channel = 0; channel < hdr6->channels; channel++) {
                    memcpy(&Y, buf, 4); buf += 4;
                    if((Y >> 30) & 3) { free(ret); ret = NULL; goto cleanup; }
                    Y &= 0x3fffffff;
                    l = (Y + 7) >> 3;
                    if(!(ptr = (uint8_t*)malloc(l))) { free(ret); ret = NULL; goto cleanup; }
                    memcpy(ptr, buf, l); buf += l;
                    tlg6_decode_golomb_values((uint8_t*)pixbuf + channel, pixel_count, ptr);
                    free(ptr);
                }
                ptr = ftypes + (y >> 3) * x_block_count;
                skip = (ylim - y) * 8;
                for(Y = y; Y < ylim; Y++) {
                    curr_line = &ret[Y * hdr6->width];
                    j = ((ylim - Y -1) - (Y - y)); k = (Y & 1) ^ 1;
                    if(main_count) {
                        l = ((hdr6->width < 8) ? hdr6->width : 8) * (Y - y);
                        tlg6_decode_line(prev_line, curr_line, hdr6->width, 0, main_count, ptr, skip, pixbuf + l,
                            hdr6->channels == 3 ? 0xff000000 : 0, j, k);
                    }
                    if(main_count != x_block_count) {
                        x = fraction; if(x > 8) x = 8;
                        l = x * (Y - y);
                        tlg6_decode_line(prev_line, curr_line, hdr6->width, main_count, x_block_count, ptr, skip, pixbuf + l,
                            hdr6->channels == 3 ? 0xff000000 : 0, j, k);
                    }
                    prev_line = curr_line;
                }
            }
        }
    }
    if(ret && *w && *h)
        for(i = 0, j = (*w) * (*h), ptr = (uint8_t*)ret; i < j; i++, ptr += 4) { a = ptr[0]; ptr[0] = ptr[2]; ptr[2] = a; }
cleanup:
    for(channel = 0; channel < 4; channel++)
        if(blkdata[channel]) free(blkdata[channel]);
    if(ftypes) free(ftypes);
    if(zero_line) free(zero_line);
    if(pixbuf) free(pixbuf);
    return ret;
}

#endif

#endif
